#
#   HarmonicOscillator1D.py
#
#   A Library to define the system's parameters, and its wavefunction.
#
#   It might be faster to write this in C or Fortran, because it gets
#   called over and over... but this is fast enough for our uses.
#

#
#   Copyright (C) 2012
#
#   This program is free software: you can redistribute it and/or modify
#   it under the terms of the GNU General Public License as published by
#   the Free Software Foundation, either version 3 of the License, or
#   (at your option) any later version.
#
#   This program is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU General Public License for more details.
#
#   You should have received a copy of the GNU General Public License
#   along with this program.  If not, see <http://www.gnu.org/licenses/>.
#


from math import sqrt
from math import sin
from cmath import pi
from cmath import exp
from cmath import phase
from math import factorial

# Problem setup: harmonic potential of `frequency' omega

hbar = 1.0
mass = 1.0
omega = 1.0

def V(x):
    return 0.5 * mass * omega**2 * x**2

# -----------------------------------------------------------------------------------------------

# Start with coherent states.

def make_Coherent_State(alpha): # Give the constant alpha(0), not alpha(t) !!!
    # useful combinations of coefficients
    norm1 = (mass*omega/pi/hbar)**0.25
    norm2 = sqrt(2*hbar/mass/omega)
    norm3 = -mass*omega/2/hbar
    norm4 = sqrt(2*mass*omega*hbar)
    # Heisenberg Picture alpha(t)
    alpha_t = lambda t: alpha*exp((0-1j)*omega*t)
    # construct the amplitude
    x_hat_average = lambda t: norm2*(alpha_t(t)).real
    R = lambda t,x : norm1 * exp(norm3 * (x-x_hat_average(t))**2)
    # construct the phase
    p_hat_average = lambda t: norm4 * alpha_t(t).imag
    delta = lambda t: abs(alpha)**2 * sin(2*omega*t - 2 * phase(alpha)) / 2.0
    S = lambda t,x : p_hat_average(t)*x/hbar - omega*t/2.0+delta(t)
    # finally, return the coherent state as a function.
    return lambda t,x : R(t,x)*exp((0+1j)*S(t,x))

# for example...
CPsi = make_Coherent_State(0-1j)

#---------------------------------------------------------------------------------------------------

# We will define many eigenfunctions, which all have similar expressions.
# Rather than have long lists of definitions
#
#   Psi1 = ...
#   Psi2 = ...
#       ...
#   PsiN = ...
#
# we use a common metaprogramming trick: generate the definitions as strings and evaluate them in the desired scope.
# This trick also allows us to perform as many operations as possible at `compile time' (=> shorter runtimes).
#

def Hermite_Coeff(n,l): # Nth polynomial, L=(N-2*M)th degree in x.
    m = float(n-l)/2.0
    try:
        return factorial(n)*(-1)**m/(factorial(m)*factorial(l))*2**l
    except ValueError: # if m is negative, fractional... then the coefficient is zero.
        return 0

def Hermite_Poly(n):
    returnable_string = str(int(Hermite_Coeff(n,n))) + " * ("+str.format('{0:.17e}',sqrt(mass*omega/hbar))+"*x)**" + str(n)
    for i in range(n-2,-1,-2):
        returnable_string += " + " + str(int(Hermite_Coeff(n,i))) + " * ("+str.format('{0:.17e}',sqrt(mass*omega/hbar))+"*x)**" + str(i)
    return "("+returnable_string+")"

def make_eigenfunction_n(n):
    # This generates the definition code for the Nth eigenfunction, as a string.
    # We try to evaluate as much of the algebra as possible at `compile time', because
    # profiling had revealed that the majority of runtime was spent in these eigenfunctions.
    
    defn = "lambda t,x : "
    
    normalisation = str.format('{0:.17e}',eval("1.0/sqrt(2.0**"+str(n)+" * factorial("+str(n)+"))*("+str((mass*omega/(hbar*pi)))+")**0.25"))
    gauss = str.format('{0:.17e}',exp(-mass*omega/hbar/2.0).real)+"**(x**2)"
    time = "exp("+str((0-1j)*omega*(n+0.5))+"*t)"

    # Now, put everything together:
    
    # Given strings "a","b","c",... this function returns the string "a*b*c*..."
    str_product = lambda a,b,*args: str_product(a+"*"+b,*args) if args else a+"*"+b
    
    todo = defn + str_product(normalisation, gauss, Hermite_Poly(n), time)
            
    return todo

    # The `To Do' string contains the definition of the eigenstate:
    # Psi_N = lambda t,x : Normalisation * Gaussian * Hermite_Polynomial * Time_Dependence

# We can now loop over N and produce all the definitions by executing the generated code:

for i in range(0,151):
    # Beyond n=150, the normalisation constant becomes too big to convert from a long integer data type to a float data type.
    #print(make_eigenfunction_n(i))
    exec("Psi"+str(i)+"="+make_eigenfunction_n(i)) # Defines Psi1, Psi2, ...

#---------------------------------------------------------------------------------------------------

# Given eigenfunctions, we can now define any system wavefunction as a sum:

def make_Coherent_Superposition(alpha): # Give the constant alpha(0), not alpha(t) !!!
    sum = "0"
    for n in range(21):
        sum = str(alpha**n /sqrt(factorial(n)))+" * Psi"+str(n)+"(t,x) + " + sum
    evaluatable = "lambda t,x : "+ str(exp(-alpha.conjugate()*alpha/2.0).real) + " * ("+sum+")"
    return eval(evaluatable)

# for example
CPsi_2 = make_Coherent_Superposition(0-1j) # 0-1j

def make_Massive_Psi(how_many):
    # in equal superposition...
    sum = "0"
    for i in range(how_many):
        sum = "Psi"+str(i)+"(t,x) + " + sum
    return eval("lambda t,x : "+str(1.0/sqrt(how_many))+" * ("+sum+")")

# for example...
massive_Psi = make_Massive_Psi(8)

invsqrt2 = 1.0/sqrt(2.0)

def Psi(t,x): # the superposition (equal contributions)
    #return invsqrt2 * (Psi1(t,x) + Psi2(t,x))
    #return Psi9(t,x)
    return CPsi(t,x)
    #return massive_Psi(t,x)

#---------------------------------------------------------------------------------------------------

#                                       TESTS

#---------------------------------------------------------------------------------------------------


if __name__ == "__main__":
    print("\nRunning Diagnostic Tests...\n")

    print("Potential Strength:")
    print("\omega = "+str(omega))
    
    print("\nCheck that the normalisation of the exported wavefunction is 1.0:")

    times = 5000
    # For the integration limits, +/- 20 is close enough to infinity
    comb = [i/float(times)/2.0 for i in range(-40*times,40*times+1)]
    time = [0.17 for i in range(-40*times,40*times+1)] # Any time should work fine

    a = map(Psi,time,comb)
    delta_comb = comb[7]-comb[6]
    b = [abs(psi)**2 for psi in a]
    from numpy import sum
    print(str.format('1.0 + {0:.5e}',1 - sum(b) * delta_comb))
        
    print("\nVerify Eigenstate Orthonormality:\n")

    diffs = []

    for i in range(4): # checked up to n=20 already :-)
        print(i)
        a = eval("map(Psi"+str(i)+",time,comb)")
        for j in range(i+1):
            b = eval("map(Psi"+str(j)+",time,comb)")
            c = map(lambda x,y: x.conjugate()*y, a, b)
            d = sum(c).real * delta_comb
            if d > 0.5:
                d = 1 - d
            diffs.append(abs(d))

    print(max(diffs))

    print("\nVerify coherent states:\n")
    print(max([abs(CPsi(0.0,x/10.0)-CPsi_2(0.0,x/10.0)) for x in range(-100,101)]))

    print("\nDone.\n")